package cn.iocoder.yudao.module.dify.config;

import io.github.imfangs.dify.client.DifyClientFactory;
import io.github.imfangs.dify.client.DifyDatasetsClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class DifyKnowledgeBaseManager {

    @Value("${spring.dify.baseUrl}")
    private String baseUrl;

    private final Map<String, String> datasetKeys = new HashMap<>();
    private final Map<String, DifyDatasetsClient> datasetsClientMap = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        // 可以从配置文件加载工作流键值对
        loadWorkflowKeys();
        initializeClients();
    }

    private void loadWorkflowKeys() {
        // 从配置或数据库加载工作流配置
        datasetKeys.put("batchUpload", "dataset-XknqKAVHyTt4V0BUcFaPYLK0");//批量上传知识库
    }

    private void initializeClients() {
        datasetKeys.forEach((name, key) ->
                datasetsClientMap.put(name, DifyClientFactory.createDatasetsClient(baseUrl, key))
        );
    }

    public DifyDatasetsClient getClient(String datasetKey) {
        return datasetsClientMap.get(datasetKey);
    }

    public Set<String> getAvailableDatasets() {
        return datasetsClientMap.keySet();
    }
}

